package com.example.demo.filters;

import com.example.demo.domain.ResultMessage;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RateLimiterRegistry;
import io.vavr.control.Try;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;

/**
 * 多服务限流过滤器
 */
@Component
public class RateLimiterFilter implements GatewayFilter, Ordered {

    @Autowired
    private RateLimiterRegistry rateLimiterRegistry;
    
    // 服务路径前缀与限流器名称的映射
    private static final Map<String, String> SERVICE_LIMITER_MAP = new HashMap<>();
    
    // 初始化服务与限流器的映射关系
    static {
        SERVICE_LIMITER_MAP.put("/user/", "user");
        SERVICE_LIMITER_MAP.put("/product/", "product");
        SERVICE_LIMITER_MAP.put("/fund/", "fund");
        // 可以根据需要添加更多服务映射
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 获取请求路径
        String path = exchange.getRequest().getURI().getPath();
        
        // 确定使用哪个限流器
        String limiterName = determineRateLimiterName(path);
        
        // 应用限流逻辑
        ResultMessage result = rateLimit(limiterName);
        
        if (!result.isSuccess()) {
            ServerHttpResponse response = exchange.getResponse();
            response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            response.getHeaders().setContentType(MediaType.APPLICATION_JSON_UTF8);
            String body = toJson(result);
            DataBuffer dataBuffer = response.bufferFactory().wrap(body.getBytes(StandardCharsets.UTF_8));
            return response.writeWith(Mono.just(dataBuffer));
        }
        return chain.filter(exchange);
    }

    /**
     * 根据请求路径确定使用哪个限流器
     * @param path 请求路径
     * @return 限流器名称
     */
    private String determineRateLimiterName(String path) {
        return SERVICE_LIMITER_MAP.entrySet().stream()
                .filter(entry -> path.startsWith(entry.getKey()))
                .map(Map.Entry::getValue)
                .findFirst()
                .orElse("default"); // 如果没有匹配的服务，使用默认限流器
    }

    /**
     * 应用限流逻辑
     * @param limiterName 限流器名称
     * @return 限流结果
     */
    private ResultMessage rateLimit(String limiterName) {
        // 获取指定的限流器，如果不存在则使用默认限流器
        RateLimiter serviceLimiter;
        try {
            serviceLimiter = rateLimiterRegistry.rateLimiter(limiterName);
        } catch (Exception e) {
            // 如果指定的限流器不存在，使用默认限流器
            serviceLimiter = rateLimiterRegistry.rateLimiter("default");
        }
        
        Callable<ResultMessage> call = RateLimiter.decorateCallable(serviceLimiter, 
                () -> new ResultMessage(true, "通过"));
        Try<ResultMessage> tryResult = Try.of(call::call)
                .recover(ex -> new ResultMessage(false, "请求过多，服务[" + limiterName + "]已限流"));

        return tryResult.get();
    }

    @Override
    public int getOrder() {
        return -1;
    }

    private String toJson(ResultMessage result) {
        ObjectMapper mapper = new ObjectMapper();
        String body = null;
        try {
            body = mapper.writeValueAsString(result);
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        return body;
    }
}
